import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from LSTM_one.LSTM.dataset_create import create_dataset
from LSTM_one.LSTM.lstm import LSTM

def train_and_save(province):
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 图中字体改为黑体以兼容中文
    plt.rcParams['axes.unicode_minus'] = False  # 负号显示的问题
    # 检查省份信息是否为 None，如果是则设置默认值
    if province is None:
        province = 'default_province'

    # 构建数据路径
    data_path = 'data/2206_2404新能源汽车总体销量数据_' + province + '.csv'

    # 打印数据路径
    print(data_path)
    data = pd.read_csv(data_path, encoding='latin1')

    # 归一化处理
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data['sale_count'].values.reshape(-1, 1))

    # 创建训练集和测试集的有监督学习数据
    look_back = 6
    train_X, train_Y = create_dataset(scaled_data, look_back)
    # 将数据转换为张量
    train_X = torch.from_numpy(train_X).type(torch.Tensor)
    train_Y = torch.from_numpy(train_Y).type(torch.Tensor)
    # 定义模型参数
    input_size = 1
    hidden_size = 120
    num_layers = 3
    output_size = 1
    learning_rate = 0.001
    num_epochs = 800

    # 定义模型、损失函数和优化器
    lstm = LSTM(input_size, hidden_size, num_layers, output_size)
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(lstm.parameters(), lr=learning_rate)

    # 训练模型
    train_loss = []
    for epoch in range(num_epochs):
        outputs = lstm(train_X)
        optimizer.zero_grad()
        loss = criterion(outputs, train_Y)
        loss.backward()
        optimizer.step()
        print('Epoch [{}/{}], Train Loss: {:.4f}'.format(epoch + 1, num_epochs, loss.item()))
        train_loss.append(loss.item())

    # 反归一化
    train_predict = lstm(train_X)
    train_predict = scaler.inverse_transform(train_predict.detach().numpy())
    train_predict = np.round(train_predict)
    train_Y = scaler.inverse_transform(train_Y.detach().numpy())

    train_real = np.array(data['sale_count'][look_back:])
    train_predict = np.array(train_predict)
    mape_, mean_mape = mape(train_real, train_predict)
    print('平均相对误差:', mean_mape)

    accuracy_, accuracy = accurecy(train_real, train_predict)
    print('accuracy:', accuracy)
    # 模型存储
    # 检查省份信息是否为 None，如果是则设置默认值
    if province is None:
        province = 'default_province'
    save_path = 'model/lstm_model_' + province + '.pth'
    torch.save(lstm.state_dict(), save_path)
    # # 绘制预测结果图
    plt.ion()
    plt.plot(pd.date_range(start='2023-01', end='2024-05', freq='M'), train_real, label='真实值')
    plt.plot(pd.date_range(start='2023-01', end='2024-05', freq='M'), train_predict, label='预测值')
    plt.xlabel('时间')
    plt.ylabel('销量')
    title = province + '省LSTM模型实际值与预测值比较'
    plt.title(title)
    plt.legend()
    plt.show()
    plt.close(1)

    # 存储loss值
    if province is None:
        province = 'default_province'
    file_path = 'loss/loss_' + province + '.txt'
    with open(file_path, 'w') as file:
        for epoch in range(num_epochs):
            file.write('Epoch [{}/{}], Train Loss: {:.4f}\n'.format(epoch + 1, num_epochs, train_loss[epoch]))
        file.write('平均相对误差:{:.8f}\n'.format(mean_mape))
        file.write('accuracy:{:.8f}\n'.format(accuracy))

    # 模型损失图
    plt.plot(train_loss, color='darkorange')
    plt.xlabel('训练次数')
    plt.ylabel('训练损失')
    plt.title('LSTM模型损失函数图')
    plt.show()
    plt.close(1)

# 计算MAPE
def mape(y_true, y_pred):
    mape=[]
    for i in range(len(y_true)):
        mape_=np.abs((y_true[i] - y_pred[i]) / y_true[i])
        mape.append(mape_)
    return np.array(mape),np.mean(mape)

#求准确率
def accurecy(y_true,y_pred):
    accuracy=[]
    for i in range(len(y_true)):
        accuracy_=1-np.abs((y_true[i]-y_pred[i])/y_true[i])
        accuracy.append(accuracy_)
    return np.array(accuracy),np.mean(accuracy)



